package ai.koog.prompt.executor.clients.anthropic

import ai.koog.prompt.executor.clients.InternalLLMClientApi
import ai.koog.prompt.executor.clients.serialization.AdditionalPropertiesFlatteningSerializer
import kotlinx.serialization.EncodeDefault
import kotlinx.serialization.EncodeDefault.Mode.ALWAYS
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonElement
import kotlinx.serialization.json.JsonObject

/**
 * Represents a request for an Anthropic message-based interaction.
 *
 * This data class is used to configure the parameters for an Anthropic message request,
 * including the model, messages, and additional settings for generation behavior.
 *
 * @property model The identifier of the Anthropic model to be used for processing the request.
 * @property messages A list of messages constituting the dialogue. Each message contains a role and corresponding content.
 * @property maxTokens The maximum number of tokens to generate in the response. Defaults to 2048.
 * @property temperature The sampling temperature for adjusting randomness in generation. Higher values increase randomness.
 * @property system An optional list of system-level messages defining additional context or instructions.
 * @property tools An optional list of tools that the Anthropic model can utilize during processing.
 * @property stream Whether responses should be returned as a stream. Defaults to false.
 * @property toolChoice An optional specification for tool selection during processing, which may define automatic, specific, or no tool usage.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicMessageRequest(
    val model: String,
    val messages: List<AnthropicMessage>,
    @SerialName("max_tokens")
    @EncodeDefault
    val maxTokens: Int = MAX_TOKENS_DEFAULT,
    val temperature: Double? = null,
    val system: List<SystemAnthropicMessage>? = null,
    val tools: List<AnthropicTool>? = null,
    @EncodeDefault
    val stream: Boolean = false,
    @SerialName("tool_choice")
    val toolChoice: AnthropicToolChoice? = null,
    val additionalProperties: Map<String, JsonElement>? = null,
) {
    init {
        require(maxTokens > 0) { "maxTokens must be greater than 0, but was $maxTokens" }
        if (temperature != null) {
            require(temperature >= 0) { "temperature must be greater than 0, but was $temperature" }
        }
    }

    /**
     * Companion object with default values for request
     */
    public companion object {
        /**
         * Default max tokens
         */
        public const val MAX_TOKENS_DEFAULT: Int = 2048
    }
}

/**
 * Represents a message within the Anthropic LLM system. This data class encapsulates
 * the role of the message and its associated content.
 *
 * @property role The role of the sender of the message, such as "user", "assistant", or "system".
 * @property content A list of content elements that form the message, where each content
 * can be of varying types like text, image, document, tool usage, or tool result.
 *
 * Note: This API is internal and subject to change or removal without notice. It is annotated
 * with [InternalLLMClientApi] to indicate its limited use case within library internals.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicMessage(
    val role: String,
    val content: List<AnthropicContent>
)

/**
 * Represents a system message with designated text and type properties
 * that can be utilized in anthropic system communication.
 *
 * @property text The content of the message.
 * @property type The type of message, defaulted to "text".
 */
@InternalLLMClientApi
@Serializable
public data class SystemAnthropicMessage(
    val text: String,
    @EncodeDefault(ALWAYS) val type: String = "text"
)

/**
 * Represents content that can be processed or generated by Anthropic systems.
 *
 * This is a sealed class with various types of content including textual data, images,
 * documents, and tool-related operations, encapsulating diverse forms of input or output content.
 * Subtypes are serialized with discriminators to support interoperability and flexibility
 * within the serialization format.
 */
@InternalLLMClientApi
@Serializable
public sealed class AnthropicContent {
    /**
     * Represents a text-based content within the AnthropicContent hierarchy.
     *
     * This class is used to encapsulate textual data. Instances of this class are serialized
     * with the discriminator "text" to identify their type in the context of polymorphic serialization.
     *
     * @property text The textual content being represented.
     */
    @Serializable
    @SerialName("text")
    public data class Text(val text: String) : AnthropicContent()

    /**
     * Represents an image content type within the AnthropicContent hierarchy.
     * This class is used to encapsulate image data, which can be sourced in different ways,
     * such as via URLs or base64-encoded strings.
     *
     * @property source The source of the image data.
     */
    @Serializable
    @SerialName("image")
    public data class Image(val source: ImageSource) : AnthropicContent()

    /**
     * Represents a document that originates from a specified source.
     * The document can have various source types, such as a URL, base64-encoded data, or plain text.
     *
     * @property source The source object describing the origin and type of the document's content.
     */
    @Serializable
    @SerialName("document")
    public data class Document(val source: DocumentSource) : AnthropicContent()

    /**
     * Represents the usage of a tool in a structured format.
     *
     * This class is part of the `AnthropicContent` sealed class hierarchy and is
     * identified by the serial name `tool_use`. It contains information about
     * the tool being used, including a unique identifier, a name to describe the
     * tool, and a JSON object for input parameters or configurations needed for the
     * tool's operation.
     *
     * @property id A unique identifier for the tool usage.
     * @property name The name of the tool being used.
     * @property input A JSON object containing input parameters for the tool's operation.
     */
    @Serializable
    @SerialName("tool_use")
    public data class ToolUse(
        val id: String,
        val name: String,
        val input: JsonObject
    ) : AnthropicContent()

    /**
     * Represents the result of a tool invocation within the Anthropic content system.
     *
     * @property toolUseId The unique identifier of the invoked tool for which this result corresponds.
     * @property content The output or result generated by the tool invocation.
     */
    @Serializable
    @SerialName("tool_result")
    public data class ToolResult(
        val toolUseId: String,
        val content: String
    ) : AnthropicContent()
}

/**
 * Represents a source of an image. This sealed interface has two implementations: one wrapping a URL and another
 * wrapping raw base64-encoded image data.
 *
 * Note: This API is internal and should not be used outside intended usage contexts. Its structure and behavior
 * may change without notice.
 */
@InternalLLMClientApi
@Serializable
public sealed interface ImageSource {
    /**
     * Represents an image source defined by a URL.
     *
     * This class is a part of the `ImageSource` sealed interface and specifies
     * that the image is sourced from a web resource identified by the provided URL.
     *
     * @property url The URL pointing to the image resource.
     */
    @Serializable
    @SerialName("url")
    public data class Url(val url: String) : ImageSource

    /**
     * Represents image data encoded in Base64 format. This class contains the encoded
     * string representation of the image along with its media type information.
     *
     * @property data The Base64 encoded string of the image.
     * @property mediaType The media type of the image (e.g., "image/png", "image/jpeg").
     */
    @Serializable
    @SerialName("base64")
    public data class Base64(val data: String, val mediaType: String) : ImageSource
}

/**
 * Represents a sealed interface for different types of document sources.
 * It defines a contract for how documents can be represented or accessed.
 *
 * This API is internal and subject to change without notice.
 */
@InternalLLMClientApi
@Serializable
public sealed interface DocumentSource {
    /**
     * Represents a URL source for a document.
     *
     * This class specifies the URL from which a document is accessed.
     * It is part of the `DocumentSource` sealed interface and is used
     * to identify a document's location via a web address.
     *
     * @property url The URL string pointing to the document source.
     */
    @Serializable
    @SerialName("url")
    public data class Url(val url: String) : DocumentSource

    /**
     * Represents a document source that encodes data using Base64 format.
     *
     * This class provides a way to specify document content as a Base64-encoded string
     * along with its associated media type for describing the content type.
     *
     * @property data The Base64-encoded string representing the document content.
     * @property mediaType A string indicating the media type (MIME type) of the encoded document.
     */
    @Serializable
    @SerialName("base64")
    public data class Base64(val data: String, val mediaType: String) : DocumentSource

    /**
     * Represents a plain text document source, containing the text data and its associated media type.
     *
     * @property data The plain text content of the document.
     * @property mediaType The media type of the document, typically indicating the format or type of text content.
     */
    @Serializable
    @SerialName("text")
    public data class PlainText(val data: String, val mediaType: String) : DocumentSource
}

/**
 * Represents a tool definition for Anthropic integration.
 *
 * This class holds metadata about a tool including its name, description, and input schema.
 * It is used to define and describe tools in the Anthropic ecosystem, where tools are
 * entities with specific input requirements represented by the `AnthropicToolSchema`.
 *
 * @property name The unique name of the tool.
 * @property description A human-readable description of the tool's purpose or functionality.
 * @property inputSchema The schema representing the structure of the input required by the tool.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicTool(
    val name: String,
    val description: String,
    val inputSchema: AnthropicToolSchema
)

/**
 * Represents a schema definition for an Anthropic tool utilized in LLM clients. This data class
 * defines the structure expected for tools, including the type of schema, properties, and required fields.
 *
 * This API is internal and should not be used outside its intended scope, as it might be subject
 * to changes or removal without notice.
 *
 * @property properties A JSON object representing the properties within this schema.
 * @property required A list of property names that are mandatory within this schema.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicToolSchema(
    val properties: JsonObject,
    val required: List<String>
) {
    /**
     * The type of the schema. Always returns "object" for Anthropic tool schemas.
     */
    val type: String = "object"
}

/**
 * Represents the response structure from an Anthropic API call.
 *
 * This class encapsulates essential data returned from the API, including information
 * about the response's unique identifier, type, role, content, the model used, and additional
 * metadata such as stop reason and usage statistics.
 *
 * @property id A unique identifier for the response.
 * @property type The type of the response.
 * @property role Indicates the role of the entity generating the response, such as "assistant" or "system".
 * @property content A list of content units within the response, which can represent text or tool usage instructions.
 * @property model The name or identifier of the model utilized to generate this response.
 * @property stopReason Optional. The reason why the generation process was stopped, if applicable.
 * @property usage Optional. Usage statistics for the response, including details such as token counts.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicResponse(
    val id: String,
    val type: String,
    val role: String,
    val content: List<AnthropicResponseContent>,
    val model: String,
    val stopReason: String? = null,
    val usage: AnthropicUsage? = null
)

/**
 * Represents the content of a response from Anthropic's API. This is a sealed class that encapsulates
 * different possible response types.
 *
 * Note: This API is marked as internal and should not be used directly. It may change or be removed without notice.
 */
@InternalLLMClientApi
@Serializable
public sealed class AnthropicResponseContent {
    /**
     * Represents the textual content in a response from the Anthropic language model.
     *
     * This class is a specific type of [AnthropicResponseContent], providing a structured
     * way to handle text-based responses.
     *
     * @property text The content of the textual response.
     */
    @Serializable
    @SerialName("text")
    public data class Text(val text: String) : AnthropicResponseContent()

    /**
     * Represents the usage of a tool in a response from an Anthropic system.
     *
     * @property id The unique identifier of the tool being used.
     * @property name The name of the tool being used.
     * @property input The input parameters provided to the tool, represented as a JSON object.
     */
    @Serializable
    @SerialName("tool_use")
    public data class ToolUse(
        val id: String,
        val name: String,
        val input: JsonObject
    ) : AnthropicResponseContent()
}

/**
 * Represents the usage statistics of the Anthropic LLM API.
 *
 * @property inputTokens The number of tokens sent as input to the LLM. Optional in streaming responses.
 * @property outputTokens The number of tokens received as output from the LLM. Optional in streaming responses.
 *
 * Note: This API is marked with [InternalLLMClientApi] and is intended for internal use only.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicUsage(
    val inputTokens: Int? = null,
    val outputTokens: Int? = null,
)

/**
 * Represents a response from the Anthropic stream API.
 *
 * This data class encapsulates the structure of a streamed response,
 * including its type, any delta updates to the content, and the complete message data when applicable.
 *
 * @property type The type of the response (e.g., "content_block_start", "content_block_delta", "message_delta").
 * @property index The index of the content block in streaming responses. Present in content block events.
 * @property contentBlock The content block data for "content_block_start" events, containing tool use information.
 * @property delta An optional incremental update to the message content, represented as [AnthropicStreamDelta].
 * @property message An optional complete response message, represented as [AnthropicResponse].
 * @property usage Optional usage statistics for the response, including token counts.
 * @property error Optional error information if an error occurred during streaming.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicStreamResponse(
    val type: String,
    val index: Int? = null,
    val contentBlock: AnthropicContent? = null,
    val delta: AnthropicStreamDelta? = null,
    val message: AnthropicResponse? = null,
    val usage: AnthropicUsage? = null,
    val error: AnthropicStreamError? = null
)

/**
 * Represents a delta update from the Anthropic stream response.
 *
 * This class encapsulates information regarding the type of update, optional textual content,
 * and optional tool usage data.
 *
 * @property type The type of the update provided by Anthropic.
 * @property text Optional text content associated with the delta update.
 * @property partialJson Optional partial JSON content for tool use streaming.
 * @property stopReason Optional reason why the generation process was stopped, if applicable.
 * @property toolUse Optional tool usage data associated with the delta update (deprecated).
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicStreamDelta(
    val type: String? = null,
    val text: String? = null,
    val partialJson: String? = null,
    val stopReason: String? = null,
    val toolUse: AnthropicResponseContent.ToolUse? = null
)

/**
 * Represents an error that occurred during Anthropic streaming response processing.
 *
 * This data class encapsulates error information received from the Anthropic API
 * during streaming operations, providing details about the type and nature of the error.
 *
 * @property type The type or category of the error that occurred.
 * @property message An optional descriptive message providing additional details about the error.
 * Defaults to null if no message is provided.
 */
@InternalLLMClientApi
@Serializable
public data class AnthropicStreamError(
    val type: String,
    val message: String? = null,
)

/**
 * Represents a sealed interface for different tool choices in the Anthropic client API.
 * This API is marked as internal and may change or be removed without notice.
 */
@InternalLLMClientApi
@Serializable
public sealed interface AnthropicToolChoice {
    /**
     * Represents an automatic tool choice within the AnthropicToolChoice hierarchy.
     *
     * This object indicates that the tool selection process is delegated to
     * be automatically determined rather than explicitly specified.
     *
     * It is a singleton implementation of the `AnthropicToolChoice` interface.
     */
    @Serializable
    @SerialName("auto")
    public data object Auto : AnthropicToolChoice

    /**
     * Represents a choice for using any available Anthropic tool.
     *
     * This object is part of the AnthropicToolChoice sealed interface and
     * specifies that any tool can be used without further restrictions.
     */
    @Serializable
    @SerialName("any")
    public data object Any : AnthropicToolChoice

    /**
     * Represents the absence of a tool choice in the AnthropicToolChoice hierarchy.
     *
     * This object specifies that no tool should be used.
     */
    @Serializable
    @SerialName("none")
    public data object None : AnthropicToolChoice

    /**
     * Represents a specific tool within the AnthropicToolChoice hierarchy.
     *
     * This class is a data representation of a tool, identified by its name. It extends the
     * AnthropicToolChoice sealed interface and is used to define tools by their string identifiers.
     *
     * @property name The name of the tool as a string.
     */
    @Serializable
    @SerialName("tool")
    public data class Tool(val name: String) : AnthropicToolChoice
}

internal object AnthropicMessageRequestSerializer :
    AdditionalPropertiesFlatteningSerializer<AnthropicMessageRequest>(AnthropicMessageRequest.serializer())
